import time

from scipy.io import wavfile

from common.config import conf
from common.tts import TTS
from melo.api import TTS as meloTTS


class MeloTTS(TTS):
    def __init__(self):
        super().__init__()
        config = conf()
        tts_device = config["tts_device"]
        self.speed = config["tts_speed"]
        melo_model = config["melo_model"]
        melo_config = config["melo_config"]
        self.melo = meloTTS(language='ZH', config_path=melo_config, ckpt_path=melo_model, device=tts_device)
        self.speaker_ids = self.melo.hps.data.spk2id

    def pretreatment(self):
        pass

    def text_to_voice(self, text):
        return self.melo.tts_to_file("你好，我是台湾妹子，很高兴认识你。", self.speaker_ids['ZH'], speed=self.speed)

    def text_to_voice_stream(self, text):
        pass


if __name__ == '__main__':
    melo_tts = MeloTTS()
    voice = melo_tts.text_to_voice("你好，我是台湾妹子，很高兴认识你。")
    start_time = time.time()
    voice2 = melo_tts.text_to_voice("你好，我是台湾妹子，很高兴认识你。")
    end_time = time.time()
    print(end_time - start_time)
    wavfile.write('湾妹音频1.wav', 44100, voice2)
